import os
import time
import cv2
import numpy as np
import ctypes
from ctypes import c_bool, c_double, c_ubyte, c_void_p
from gym import Env
from gym.spaces import Box, Discrete
from stable_baselines3 import DQN
from stable_baselines3.common.vec_env import DummyVecEnv
from stable_baselines3.common.callbacks import BaseCallback
import platform

# Load DLL
class Vector2(ctypes.Structure):
    _fields_ = [
        ("x", ctypes.c_double),
        ("y", ctypes.c_double),
    ]

class DinoInfo(ctypes.Structure):
    _fields_ = [
        ("position", Vector2),
        ("velocity", Vector2),
        ("collision_box_lu", Vector2),
        ("collision_box_rd", Vector2),
        ("state", ctypes.c_int),
        ("health", ctypes.c_int64),
        ("energy", ctypes.c_int64),
    ]

class FireballInfo(ctypes.Structure):
    _fields_ = [
        ("position", Vector2),
        ("velocity", Vector2),
        ("collision_box_lu", Vector2),
        ("collision_box_rd", Vector2),
    ]

class PterosaurInfo(ctypes.Structure):
    _fields_ = [
        ("position", Vector2),
        ("velocity", Vector2),
        ("collision_box_lu", Vector2),
        ("collision_box_rd", Vector2),
    ]

class CactusInfo(ctypes.Structure):
    _fields_ = [
        ("position", Vector2),
        ("velocity", Vector2),
        ("collision_box_lu", Vector2),
        ("collision_box_rd", Vector2),
    ]

class MeteoriteInfo(ctypes.Structure):
    _fields_ = [
        ("position", Vector2),
        ("velocity", Vector2),
        ("collision_box_lu", Vector2),
        ("collision_box_rd", Vector2),
    ]

class GameInfo(ctypes.Structure):
    _fields_ = [
        ("dinosaur", ctypes.POINTER(DinoInfo)),
        ("fireballs", ctypes.POINTER(FireballInfo)),
        ("fireball_count", ctypes.c_uint64),
        ("pterosaurs", ctypes.POINTER(PterosaurInfo)),
        ("pterosaur_count", ctypes.c_uint64),
        ("cactuses", ctypes.POINTER(CactusInfo)),
        ("cactus_count", ctypes.c_uint64),
        ("meteorites", ctypes.POINTER(MeteoriteInfo)),
        ("meteorite_count", ctypes.c_uint64),
        ("flag_fail", ctypes.c_bool),
        ("flag_pause", ctypes.c_bool),
        ("score", ctypes.c_uint64)
    ]

class MiniDinoGame:
    def __init__(self, to_render=True, with_window=True, to_extract_game_info=True, lib_path=None):
        system_type = platform.system()
        if lib_path is None:
            if system_type == "Windows":
                libfile_name = "mini_dino.dll"
            elif system_type == "Linux":
                libfile_name = "libmini_dino.so"
            else:
                raise RuntimeError("Unsupported os.")

        # lib_path = "D:/zhulibo/university/大三/智能系统控制/chrome-dinosaur/PyMini-Dino/mini_dino.dll"
        lib_path = os.path.join(os.path.dirname(__file__), libfile_name)
        self.lib = ctypes.cdll.LoadLibrary(lib_path)

        # create_game(to_render: bool, with_window: bool, to_extract_game_info: bool) -> *mut GameApp
        self.lib.create_game.argtypes = [c_bool, c_bool, c_bool]
        self.lib.create_game.restype  = c_void_p

        self.lib.initialize.argtypes = [c_void_p]
        self.lib.initialize.restype = None

        self.lib.run.argtypes = [c_void_p]
        self.lib.run.restype = None

        self.lib.step.argtypes = [c_void_p, c_void_p, c_double]
        self.lib.step.restype = c_bool

        self.lib.get_game_info_ptr.argtypes = [c_void_p]
        self.lib.get_game_info_ptr.restype = ctypes.POINTER(GameInfo)

        self.lib.get_buffer_ptr.argtypes = [c_void_p]
        self.lib.get_buffer_ptr.restype = ctypes.POINTER(c_ubyte)

        self.lib.reset.argtypes = [c_void_p]
        self.lib.reset.restype = None

        self.lib.clean_up.argtypes = [c_void_p]
        self.lib.clean_up.restype = None

        self.game_ptr = self.lib.create_game(to_render, with_window, to_extract_game_info)
        if not self.game_ptr:
            raise RuntimeError("Failed to init_game, got null pointer!")

        self.is_initialized = False

        self.round_living_timer = 0.0

    def initialize(self):
        if not self.is_initialized:
            self.lib.initialize(self.game_ptr)
            self.is_initialized = True

    def run(self):
        if not self.is_initialized:
            self.initialize()

        # Game loop
        while True:
            game_info = self.get_game_info()  # Get current game info
            if game_info.flag_fail:
                print("Game Over!")
                print(f"Dino has lived for {self.round_living_timer} secs")
                self.reset()  # Reset the game after failure
                keyboard_dict = {25: True}  # 使用正确扫描码
                keyboard_ptr = self.create_keyboard_state(keyboard_dict)
                self.lib.step(self.game_ptr, keyboard_ptr, 0.0167)

            elif game_info.flag_pause:
                print("Game Paused!")
                # time.sleep(1)  # Wait for unpause
                keyboard_dict = {25: True}  # 使用正确扫描码
                keyboard_ptr = self.create_keyboard_state(keyboard_dict)
                self.lib.step(self.game_ptr, keyboard_ptr, 0.0167)

            # Proceed with normal game step if not paused or failed
            self.lib.step(self.game_ptr, None, 0.0167)
            # time.sleep(0.01)  # Sleep to avoid 100% CPU usage

    def step(self, keyboard_ptr, delta_time: float):
        # Check if the game failed or paused and handle accordingly
        self.round_living_timer += delta_time
        game_info = self.get_game_info()
        if game_info.flag_fail:
            print("Game Over!")
            keyboard_dict = {27: True}  # 使用正确扫描码
            keyboard_ptr = self.create_keyboard_state(keyboard_dict)
            self.lib.step(self.game_ptr, keyboard_ptr, 0.0167)
            return False  # Stop game
        elif game_info.flag_pause:
            print("Game Paused!")
            keyboard_dict = {25: True}  # 使用正确扫描码
            keyboard_ptr = self.create_keyboard_state(keyboard_dict)
            self.lib.step(self.game_ptr, keyboard_ptr, 0.0167)
            return False  # Game paused

        return self.lib.step(self.game_ptr, keyboard_ptr, delta_time)

    def reset(self):
        self.lib.reset(self.game_ptr)
        self.round_living_timer = 0.0

    def get_game_info(self):
        info_ptr = self.lib.get_game_info_ptr(self.game_ptr)
        if not info_ptr:
            print("Failed to get game info")
            return None
        return info_ptr.contents

    def get_game_frame(self):
        """
        The returned image is stored in BGRA format,
        which is common in OpenCV. Here we will ignore the alpha channel,
        resize the image to a smaller size, and return it as RGB.
        """
        buf_ptr = self.lib.get_buffer_ptr(self.game_ptr)
        width, height = 1000, 600
        buffer = np.ctypeslib.as_array(buf_ptr, shape=(width * height * 4,))  # 4 channels (BGRA)
        image_data = buffer.reshape((height, width, 4))  # BGRA format
        image_bgra = image_data.astype(np.uint8)
    
        # Ignore the alpha channel (last channel), keep only RGB
        image_rgb = cv2.cvtColor(image_bgra, cv2.COLOR_BGRA2RGB)
    
        # Resize the image to 224x224 (or smaller)
        resized_image = cv2.resize(image_rgb, (224, 224))
    
        return resized_image  # Now it's (224, 224, 3)


    def create_keyboard_state(self, keyboard_dict):
        keyboard_state = (ctypes.c_ubyte * 108)()
        for key, pressed in keyboard_dict.items():
            if 0 <= key < 108:
                keyboard_state[key] = 1 if pressed else 0
        return ctypes.cast(keyboard_state, ctypes.POINTER(ctypes.c_ubyte))

    def close(self):
        self.lib.clean_up(self.game_ptr)

    def __del__(self):
        self.close()


class DinoGameEnv(Env):
    def __init__(self, to_render=True, with_window=True, to_extract_game_info=True, lib_path=None):
        super().__init__()
        self.game = MiniDinoGame(to_render=to_render, with_window=with_window, to_extract_game_info=to_extract_game_info, lib_path=lib_path)
        self.game.initialize()

        self.observation_space = Box(low=0, high=255, shape=(224, 224, 3), dtype=np.uint8)
        self.action_space = Discrete(7)  # w, a, s, d, j, k ，no_op(7个操作)

        self.last_score = 0
        self.last_healthy = 0
        self.last_energy = 0
        # 定义每个动作的持续时间（单位：秒）
        self.action_durations = {
            0: 2.3,  # w
            1: 0.7,  # a
            2: 1.2,  # s
            3: 0.7,  # d
            4: 0.00001,  # j
            5: 0.00001,   # k
            6: 0.7  # z(表示无操作)
        }

    def reset(self):
        self.game.reset()
        # time.sleep(1)
        return self.get_observation()

    def step(self, action):
        keyboard_state = self.create_keyboard_state(action)

        # 获取当前动作的持续时间
        duration = self.action_durations.get(action, 1.0)  # 默认持续1秒

        # 执行动作并持续指定时间
        action_timer = 0.0
        while action_timer < duration:  # 持续指定时间
            self.game.step(keyboard_state, 0.0167)
            action_timer += 0.0167

        game_info = self.game.get_game_info()
        done = game_info.flag_fail
        reward = self.calculate_reward(game_info)
        observation = self.get_observation()

        return observation, reward, done, {}

    def render(self):
        frame = self.game.get_game_frame()
        cv2.imshow('Dino Game', frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            self.close()

    def close(self):
        cv2.destroyAllWindows()

    def get_observation(self):
        return self.game.get_game_frame()

    def create_keyboard_state(self, action):
        keyboard_dict = {
            0: {32: True},  # w
            1: {10: True},  # a
            2: {28: True},  # s
            3: {13: True},  # d
            4: {19: True},  # j
            5: {20: True},  # k
            6: {35: True}   
        }
        state = {key: False for key in range(108)}
        if action in keyboard_dict:
            for key, pressed in keyboard_dict[action].items():
                state[key] = pressed
        return self.game.create_keyboard_state(state)

    def calculate_reward(self, game_info):
        reward = 0
        dino = game_info.dinosaur.contents if game_info.dinosaur else None  # Dereference the pointer
        
        # 基础存活奖励（鼓励持续生存）
        reward += 1
        
        if dino:
            health_change = dino.health - self.last_healthy
            reward += health_change * 5  # 每损失1点生命-5
            
            energy_usage = self.last_energy - dino.energy
            reward -= energy_usage * 1.5  # 每次JK消耗2能量对应-3惩罚
            score_change = game_info.score - self.last_score
            reward += score_change * 0.2
            
            velocity_reward = abs(dino.velocity.x * 0.15)
            reward += velocity_reward
        if game_info.flag_fail:
            reward -= 20

        self.last_score = game_info.score
        self.last_healthy = dino.health
        self.last_energy = dino.energy
        return reward
    
class TrainAndLoggingCallback(BaseCallback):
    def __init__(self, check_freq, save_path, verbose=1):
        super(TrainAndLoggingCallback, self).__init__(verbose)
        self.check_freq = check_freq
        self.save_path = save_path

    def _init_callback(self):
        if self.save_path is not None:
            os.makedirs(self.save_path, exist_ok=True)

    def _on_step(self):
        if self.n_calls % self.check_freq == 0:
            model_path = os.path.join(self.save_path, 'best_model_{}'.format(self.n_calls))
            self.model.save(model_path)

        return True
    

# Environment setup
CHECKPOINT_DIR = './train/'
LOG_DIR = './logs/'

env = DummyVecEnv([lambda: DinoGameEnv(with_window=False)])

# Create and train model
model = DQN('CnnPolicy', env, tensorboard_log=LOG_DIR, verbose=1, buffer_size=4000, learning_starts=800)

callback = TrainAndLoggingCallback(check_freq=5000, save_path=CHECKPOINT_DIR)
model.learn(total_timesteps=100000, callback=callback)

# 保存模型
model.save('best_model')